import time
import json
import requests
from selenium import webdriver
from selenium.webdriver.common.by import By
from datetime import datetime, time as datetime_time

STATION_DATA_FILE = "stations-g.json"

DEFAULT_DEPARTURE = "杭州"
DEFAULT_ARRIVAL = "银川"
DEFAULT_DATE = time.strftime("%Y-%m-%d", time.localtime(time.time() + 3 * 24 * 60 * 60))

def init_driver():
	driver = webdriver.Edge()
	driver.get(
		f'https://kyfw.12306.cn/otn/leftTicket/init?linktypeid=dc&fs=%E6%9D%AD%E5%B7%9E,HZH&ts=%E9%93%B6%E5%B7%9D,YIJ&date={DEFAULT_DATE}&flag=N,N,Y')
	return driver
	# return a usable driver with cookie & live session


def get_train_data(start_city=DEFAULT_DEPARTURE, arrive_city=DEFAULT_ARRIVAL, start_date=DEFAULT_DATE, driver=None):
	# Load the city data from `CITY_DATA_FILE`
	with open(STATION_DATA_FILE, "r", encoding="utf-8") as f:
		city_data = json.load(f)

	assert start_city in city_data, f"Invalid start city: {start_city}"
	assert arrive_city in city_data, f"Invalid arrival city: {arrive_city}"
	assert start_date >= datetime.now().strftime("%Y-%m-%d"), f"Invalid start date: {start_date}"
	assert driver is not None, "Driver is not initialized"

	result = []
	# Concatenate the URL
	url = f'https://kyfw.12306.cn/otn/leftTicket/queryG?leftTicketDTO.train_date={start_date}&leftTicketDTO.from_station={city_data[start_city]}&leftTicketDTO.to_station={city_data[arrive_city]}&purpose_codes=ADULT'

	try:
		driver.get(url)
		body = driver.find_elements(By.TAG_NAME, "body")[0].get_attribute('innerHTML')
		# get index of "[" just behind "result"
		start = body.find("[", body.find("\"result\"")) + 1
		# get index of "]", which is the end of the JSON data
		end = body.find("]", start)
		# extract the JSON data
		data = body[start:end] # data is a string, without the outermost "[" and "]"
		data = data.split(",")
		for i in range(0, len(data), 2):
			result.append(data[i][1:-2])
	except (requests.RequestException, KeyError):
		print("Failed to get train information, please check the network or URL address!")
		exit(1)
	return result


def main():
	driver = init_driver()
	# result = getSchedule(driver=driver)
	while True:
		start_city = input("Departure station: ")
		if start_city == "q":
			break
		arrive_city = input("Arrival station: ")
		start_date = input("Departure date(YYYY-MM-DD): ")
		result = get_train_data(start_city, arrive_city, start_date, driver=driver)

		# Load the city data from "city.json"
		with open(STATION_DATA_FILE, "r", encoding="utf-8") as f:
			city_data = json.load(f)
		for train in result:
			train_info = train.split("|")
			train_no = train_info[3]
			# train_source = train_info[4]
			# train_terminal = train_info[5]
			train_dept = train_info[6]
			train_arrv = train_info[7]
			train_dept_time = train_info[8]
			train_arrv_time = train_info[9]
			train_duration = train_info[10]
			# Get the city name from the city code reversely
			# If not found, return the original city code
			try:
				train_dept_name = list(city_data.keys())[list(city_data.values()).index(train_dept)]
			except ValueError:
				train_dept_name = train_dept
			try:
				train_arrv_name = list(city_data.keys())[list(city_data.values()).index(train_arrv)]
			except ValueError:
				train_arrv_name = train_arrv
			# Format the output
			train_no = train_no.ljust(5)
			train_dept_name = train_dept_name.ljust(10)
			train_arrv_name = train_arrv_name.ljust(10)
			print(f"{train_no} {train_dept_name} {train_arrv_name} {train_dept_time} {train_arrv_time} {train_duration}")

if __name__ == '__main__':
	main()
